
import torch
import torch.nn as nn
import torch
from torch import random
import torch.nn as nn
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.nn.parameter import Parameter
import torch  
import torch.nn as nn  
import torch.nn as nn  
from torchvision.models import resnet18  
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms, datasets, models
import time
import copy
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
import pickle

from typing import List

# -*- coding: utf-8 -*-


import torch
import torchvision


def freeze_resnet50_layers(
    model: torch.nn.Module, freeze_layer_count: int = 7
) -> torch.nn.Module:
    """Freeze pretrained weights in some layers of resnet50 model

    Args:
        model (torch.nn.Module): Resnet50 nn module with pretrained weights
            from torchvision models
        freeze_layer_count (int, optional): Number of child layers to be frozen.
            Defaults to 7.

    Returns:
        torch.nn.Module: Resnet50 nn module with frozen weights in some layers
    """
    layer_count = 0
    for child in model.children():
        layer_count += 1
        if layer_count <= freeze_layer_count:
            for param in child.parameters():
                param.requires_grad = False
    return model


def freeze_efficientnetv2_stages(
    model: torch.nn.Module, freeze_stage_count: int = 3
) -> torch.nn.Module:
    """Frezee pretrained weights in some stages of efficientnet_v2 model

    Args:
        model (torch.nn.Module): Efficientnetv2_s nn module with pretrained weights
            from torchvision models
        freeze_stage_count (int, optional): Number of conv stages to be frozen.
            Defaults to 3

    Returns:
        torch.nn.Module: Efficientnetv2_s nn module with frozen weights in some stages
    """
    frozen_stages_prefixes = [str(stage) for stage in range(freeze_stage_count + 1)]
    for nn_module in model.children():
        for name, module in nn_module.named_modules():
            if name.split(".")[0] not in frozen_stages_prefixes:
                continue
            for param in module.parameters():
                param.requires_grad = False
    return model


def get_backbone_model(
    backbone_type: str = "resnet50", frozen_layer_count: int = 7
) -> torch.nn.Module:
    """Get backbone model from torchvision models with some frozen pretrained layers
    and remove fc layer

    Args:
        backbone_type (str, optional): Type of backbone model options
            "efficientnetv2_m or resent50". Defaults to "resnet50".
        frozen_layer_count (int, optional): Number of layers or stages to be frozen.
            Defaults to 7.

    Returns:
        torch.nn.Module: Efficientnetv2_m or resnet50 model with some frozen layers
            and no fc layer
    """
    if backbone_type == "resnet50":
        backbone = torchvision.models.resnet50(
            weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1 if frozen_layer_count > 0 else None
        )
        backbone = freeze_resnet50_layers(backbone, frozen_layer_count)
    else:
        backbone = torchvision.models.efficientnet_v2_m(
            weights=torchvision.models.EfficientNet_V2_M_Weights.IMAGENET1K_V1 if frozen_layer_count > 0 else None
        )
        backbone = freeze_efficientnetv2_stages(backbone, frozen_layer_count)
    backbone = torch.nn.Sequential(*(list(backbone.children())[:-1]))
    return backbone
class MultiLabelBinaryClassifier(torch.nn.Module):
    def __init__(
        self,
        classes: List[str],
        backbone_type: str,
        frozen_layers: int = 2,
        device: torch.device = torch.device("cuda"),
    ):
        super().__init__()

        """Multi-label binary classifier class for classification problems that
            have image as input and multiple binary values as outputs
        
        Args:
            classes (List[str]): List of classes that define number of binary
                classifiers
            backbone_type (str): Type of backbone for the model
                (resnet50 or efficientnet_v2)
            frozen_layers (int): Number of layers that the pretrained weights
                will be frozen for them
            device (torch.device): Device by which the model's weights will
                be calculated
        """
        self.backbone = get_backbone_model(backbone_type, frozen_layers).to(device)
        self.binary_classifier_heads = {}
        self.feature_trans = {}
        # self.resnet = resnet18(pretrained=True) 

        in_features = 2048 if backbone_type == "resnet50" else 1280
        for i, class_ in enumerate(classes):
            self.feature_trans[class_] = torch.nn.Sequential(
                torch.nn.Dropout(p=0.4),
                torch.nn.Linear(in_features=in_features, out_features=in_features),
                torch.nn.ReLU(),
            ).to(device)

            self.binary_classifier_heads[class_] = torch.nn.Sequential(
                torch.nn.Dropout(p=0.4),
                torch.nn.Linear(in_features=in_features, out_features=256),
                torch.nn.ReLU(),
                torch.nn.Dropout(p=0.4),
                torch.nn.Linear(in_features=256, out_features=1),
                torch.nn.Sigmoid(),
                # torch.nn.Softmax(),

            ).to(device)

    def forward(self, x):
        features = torch.flatten(self.backbone(x), 1)
        trans_feature = {}
        out = {}
        for key, classifier_val in self.binary_classifier_heads.items():
            trans_feature[key] = self.feature_trans[key](features)
            out[key] = classifier_val(trans_feature[key])
        return out,trans_feature,features
# class MultiLabelBinaryClassifier(torch.nn.Module):
#     def __init__(
#         self,
#         classes: List[str],
#         backbone_type: str,
#         frozen_layers: int = 2,
#         device: torch.device = torch.device("cuda"),
#     ):
#         super().__init__()

#         """Multi-label binary classifier class for classification problems that
#             have image as input and multiple binary values as outputs
        
#         Args:
#             classes (List[str]): List of classes that define number of binary
#                 classifiers
#             backbone_type (str): Type of backbone for the model
#                 (resnet50 or efficientnet_v2)
#             frozen_layers (int): Number of layers that the pretrained weights
#                 will be frozen for them
#             device (torch.device): Device by which the model's weights will
#                 be calculated
#         """
#         self.backbone = get_backbone_model(backbone_type, frozen_layers).to(device)
#         self.binary_classifier_heads = {}
#         # self.resnet = resnet18(pretrained=True) 

#         in_features = 2048 if backbone_type == "resnet50" else 1280
#         for i, class_ in enumerate(classes):
#             self.binary_classifier_heads[class_] = torch.nn.Sequential(
#                 torch.nn.Dropout(p=0.4),
#                 torch.nn.Linear(in_features=in_features, out_features=256),
#                 torch.nn.ReLU(),
#                 torch.nn.Dropout(p=0.4),
#                 torch.nn.Linear(in_features=256, out_features=1),
#                 torch.nn.Sigmoid(),
#                 # torch.nn.Softmax(),
                
                
#             ).to(device)

#     def forward(self, x):
#         features = torch.flatten(self.backbone(x), 1)
#         out = {}
#         for key, classifier_val in self.binary_classifier_heads.items():
#             out[key] = classifier_val(features)
#         return out

